/*
  Copyright (C) 2011, Kyungjoo Kim
  
  This file is part of UHM (Unassembled HyperMatrix) sparse direct solver.
  
  All rights reserved.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.
  
  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.
  
  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  Also add information on how to contact you by electronic and paper mail.
  
  
  Kyungjoo Kim
  iamkyungjoo@gmail.com
*/
#include "petscksp.h"

#include "uhm.hxx"
#include "uhm/plugin/petsc.hxx"
 
#define UHM_ERROR_TOL 1.0e-5

int main (int argc, char **argv)
{

  // input check
  if (argc != 7) {
    printf("Try : %s [n_thread][n_parts][decomposition][blocksize][input_file]\n\n",
           argv[0]);
    printf(" ==== decomposition method ====\n\n");
    printf(" ** Symmetric Factorization **\n");
    printf("    UHM_CHOL      ( 1),   UHM_CHOL_PIV ( 2),   UHM_CHOL_INCPIV ( 3) \n");
    printf("    UHM_LDL_NOPIV ( 4),   UHM_LDL_PIV  ( 5),   UHM_LDL_INCPIV  ( 6) \n");
    printf("\n");
    printf(" ** General Structure Symmetric Factorization **\n");
    printf("    UHM_LU_NOPIV  ( 7),   UHM_LU_PIV   ( 8),   UHM_LU_INCPIV   ( 9) \n");
    printf("    UHM_QR        (10) \n");

    return 0;
  }

  int n_threads, n_parts, decomposition, blocksize;
  char *filename;

  n_threads     = atoi( (argv[1]) );
  n_parts       = atoi( (argv[2]) );
  decomposition = atoi( (argv[3]) );
  blocksize     = atoi( (argv[4]) );
  filename      = argv[5];

  double t_base, 
    t_part_mesh, t_build_tree, 
    t_decompose, 
    t_solve_1, t_solve_2, t_petsc, t_petsc_all = 0.0, 
    t_check;

  FLA_Init();

  uhm::set_num_threads    (n_threads);
  uhm::set_hier_blocksize (blocksize);

  int sym      = ( decomposition < UHM_LU_NOPIV ? UHM_SYMMETRY : UHM_UNSYMMETRY);
  int datatype = UHM_REAL, n_rhs = 1, is_schur = false;

  PetscErrorCode ierr;
  uhm::plugin::Petsc petsc = new uhm::plugin::Petsc_;

  ierr = petsc->initialize(argc, argv);CHKERRQ(ierr);
  ierr = PetscPrintf(PETSC_COMM_WORLD,
                     ">> Petsc was initialized with argc : %d\n", argc);CHKERRQ(ierr);

  uhm::Mesh m   = petsc->get_self();
  m->set_decomposition_method(decomposition);
  m->import_file( filename );
  ierr = PetscPrintf(PETSC_COMM_WORLD,
                     ">> Imported mesh from file : %s\n", filename );CHKERRQ(ierr);

  t_base       = uhm::timer();
  uhm::partition_mesh(m, n_parts);
  t_part_mesh  = uhm::timer() - t_base;
  ierr = PetscPrintf(PETSC_COMM_WORLD,
                     ">> Partitioned mesh : %lf [sec]\n", t_part_mesh );CHKERRQ(ierr);

  petsc->set_n_rhs(n_rhs);
  petsc->lock();
  
  ierr = PetscPrintf(PETSC_COMM_WORLD,
                     ">> Mesh locked\n");CHKERRQ(ierr);
  
  for (int i=0;i<petsc->get_n_children();++i) {
    uhm::Mesh c = petsc->get_child(i);

    t_base       = uhm::timer();
    uhm::build_tree(c);
    t_build_tree = uhm::timer() - t_base;
    ierr = PetscPrintf(PETSC_COMM_WORLD,
                       ">> Built tree for %d mesh : %lf\n", 
                       i, t_build_tree);CHKERRQ(ierr);
    
    c->set_symmetry(sym);
    c->lock();

    ierr = PetscPrintf(PETSC_COMM_SELF,
                       ">> rank %d, child %d, # nodes %d, # elts %d\n", 
                       petsc->get_comm_rank(), i, 
                       c->get_n_nodes(), c->get_n_elements());CHKERRQ(ierr);

    c->create_matrix_without_buffer( datatype, n_rhs );
    c->create_matrix_buffer(is_schur);
    if (decomposition < UHM_LU_NOPIV) 
      c->random_spd_matrix();
    else 
      c->random_matrix();
    c->set_rhs();
  }
  ierr = PetscPrintf(PETSC_COMM_WORLD,
                     ">> Built tree for children mesh(%d) end\n", 
                     petsc->get_n_children());CHKERRQ(ierr);



  t_base      = uhm::timer();
  petsc->decompose_local_with_free();
  t_decompose = uhm::timer() - t_base;
  ierr = PetscPrintf(PETSC_COMM_WORLD,
                     ">> Decompose children matrices : %lf\n", 
                     t_decompose);CHKERRQ(ierr);


  t_base      = uhm::timer();
  petsc->solve_1_local();
  t_solve_1   = uhm::timer() - t_base;
  ierr = PetscPrintf(PETSC_COMM_WORLD,
                     ">> Solve 1 children matrices : %lf\n",
                     t_solve_1);CHKERRQ(ierr);

  t_base      = uhm::timer();
  ierr = petsc->create_matrix();CHKERRQ(ierr);
  t_petsc     = uhm::timer() - t_base;
  t_petsc_all += t_petsc;
  ierr = PetscPrintf(PETSC_COMM_WORLD,
                     ">> Petsc create_matrix : %lf\n", 
                     t_petsc);CHKERRQ(ierr);

  ierr = petsc->view(UHM_PLUGIN_PETSC_MAT_A, PETSC_VIEWER_DRAW_WORLD);CHKERRQ(ierr);

  ierr = petsc->create_ksp();CHKERRQ(ierr);
  ierr = petsc->set_pc_on_ksp(PCBJACOBI);CHKERRQ(ierr);

  t_base      = uhm::timer();
  ierr = petsc->solve_dist();CHKERRQ(ierr);
  t_petsc     = uhm::timer() - t_base;
  t_petsc_all += t_petsc;
  ierr = PetscPrintf(PETSC_COMM_WORLD,
                     ">> Petsc solve : %lf\n", 
                     t_petsc);CHKERRQ(ierr);

  ierr = petsc->view(UHM_PLUGIN_PETSC_KSP, PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);

  ierr = petsc->check_dist();CHKERRQ(ierr);
  ierr = petsc->delete_ksp();CHKERRQ(ierr);

  t_base      = uhm::timer();
  ierr = petsc->delete_matrix();CHKERRQ(ierr);
  t_petsc     = uhm::timer() - t_base;
  t_petsc_all += t_petsc;
  ierr = PetscPrintf(PETSC_COMM_WORLD,
                     ">> Petsc delete_matrix : %lf\n", t_petsc);CHKERRQ(ierr);

  t_petsc_all += t_petsc;
  ierr = PetscPrintf(PETSC_COMM_WORLD,
                     ">> Solve Petsc : %lf\n",
                     t_petsc_all);CHKERRQ(ierr);

  t_base      = uhm::timer();
  petsc->solve_2_local();
  t_solve_2   = uhm::timer() - t_base;
  ierr = PetscPrintf(PETSC_COMM_WORLD,
                     ">> Solve 2 children matrices : %lf\n",
                     t_solve_2);CHKERRQ(ierr);

  t_base      = uhm::timer();
  petsc->check_local();
  t_check = uhm::timer() - t_base;
  ierr = PetscPrintf(PETSC_COMM_WORLD,
                     ">> Check children matrices : %lf\n",
                     t_check);CHKERRQ(ierr);

  for (int i=0;i<petsc->get_n_children();++i) {
    double res = petsc->get_child(i)->get_residual();
    ierr = PetscPrintf(PETSC_COMM_SELF,
                       ">> Rank %d, Child %d, Residual : % E\n",
                       petsc->get_comm_rank(), i, res);CHKERRQ(ierr);
  }

  petsc->unlock();
  delete petsc;

  ierr = petsc->finalize();CHKERRQ(ierr);
  ierr = PetscPrintf(PETSC_COMM_WORLD,
                     ">> Petsc was finalized\n");CHKERRQ(ierr);

  FLA_Finalize();

  return 0;
}


